package vulnerability

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"prismx_cli/core/models"
	"prismx_cli/utils/logger"
	"prismx_cli/utils/netUtils"
	"prismx_cli/utils/temp"
	"regexp"
	"strconv"
	"strings"
	"time"
)

var c = temp.NewEnvOption()
var env, _ = temp.NewEnv(&c)

// HttpVerify YAML类型利用模块
// 使用内置插件前必须初始化所遇到插件。
// 保留变量名：exploit 禁止使用，只能插件使用
func HttpVerify(target string, vlMeta models.VulMeta, payload string, isRed bool, timeout time.Duration) models.VulResult {
	var (
		request  *http.Request
		response netUtils.Result
		err      error
	)
	//----------------------------------------------
	//开始校验本次请求的响应规则

	verifyOrExploit := vlMeta.Steps.VerifySteps.Verify
	if isRed && vlMeta.Steps.ExploitSteps.Exploit != nil && payload != "" {
		verifyOrExploit = vlMeta.Steps.ExploitSteps.Exploit
	}
	verifyOrExploitJson, err := json.Marshal(verifyOrExploit)
	if err != nil {
		logger.ScanFailuref("Plug in [%s] configuration error, skipping this verification ~ error message: %s", vlMeta.Name, err.Error())
		return models.VulResult{}
	}

	var variableError int //用来统计错误的变量数量

	//替换变量
	for _, variableItem := range vlMeta.Steps.Variable {
		//修改变量
		for key, value := range variableItem {
			if strings.Contains(string(verifyOrExploitJson), fmt.Sprintf("{{%s}}", key)) {
				evaluate, err := temp.Evaluate(env, value.(string), variableItem) //如果遇到变量值设置错误的情况直接跳过该变量循环下一个
				if err != nil {
					variableError++
					logger.ScanFailuref("Error in variable setting of plug-in [%s], skipping this verification ~ error message: %s", vlMeta.Name, err.Error())
					continue
				}
				val := evaluate.Value().(string)
				verifyOrExploitJson = []byte(strings.ReplaceAll(string(verifyOrExploitJson), fmt.Sprintf("{{%s}}", key), val))
			}
		}
	}

	//替换exp的payload参数
	if strings.Contains(string(verifyOrExploitJson), "{{exploit}}") && isRed {
		verifyOrExploitJson = bytes.ReplaceAll(verifyOrExploitJson, []byte("{{exploit}}"), []byte(payload))
	}

	//如果错误次数与Variable数量相等，则不再看此次验证
	//好家伙就几个变量还全错，这你还搞nm呢
	if variableError >= len(vlMeta.Steps.Variable) && variableError != 0 {
		return models.VulResult{State: false, Response: fmt.Sprintf("Error in configuration of plug-in [%s]", vlMeta.Name)}
	}

	var verifyOrExploitTmp []models.StepMeta
	if err = json.Unmarshal(verifyOrExploitJson, &verifyOrExploitTmp); err != nil {
		logger.ScanFailuref("Error in configuration of plug-in [%s], ignoring ~ error information in this verification: %s", vlMeta.Name, err.Error())
		return models.VulResult{State: false, Response: err.Error()}
	}

	verifyOrFlag := false

	for _, item := range verifyOrExploitTmp {
		if item.Request.Params != "" {
			request, err = http.NewRequest(item.Request.Method, target+item.Request.Path, strings.NewReader(item.Request.Params))
		} else {
			request, err = http.NewRequest(item.Request.Method, target+item.Request.Path, nil)
		}
		if err != nil {
			return models.VulResult{State: false, Response: err.Error()}
		}

		//构造header
		for _, headerItem := range item.Request.Header {
			for key, value := range headerItem {
				request.Header.Set(key, value)
			}
		}

		//获取请求时间
		start := time.Now()
		//----------------------------------------------
		//http请求构造完毕
		//下面开始根据荷载进行响应数据判断
		response, err = netUtils.SendHttp(request, timeout, item.Request.Redirect)
		//如果此次请求失败且判断条件为or就继续下次循环
		//如果所有的请求都无法执行的话，那就只能执行异常
		if err != nil {
			if vlMeta.Steps.VerifySteps.Type == "and" {
				return models.VulResult{State: false, Request: response.RequestRaw, Response: err.Error()}
			}
			response.Body = []byte(err.Error())
			continue
		}

		//获取请求用时
		requestTime := time.Now().Sub(start)

		//释放资源
		if response.Other.Body != nil {
			response.Other.Body.Close()
		}
		respFlag := false
		//校验每个响应体的验证规则
		//必须code、body都要符合规则
		for _, respItem := range item.Response {
			switch respItem.Type {
			case "contains":
				if respItem.Name == "header" && strings.Contains(response.Header, respItem.Value) {
					respFlag = true
				} else if respItem.Name == "body" && strings.Contains(string(response.Body), respItem.Value) {
					respFlag = true
				} else if respItem.Name == "code" && strings.Contains(strconv.Itoa(response.Other.StatusCode), respItem.Value) {
					respFlag = true
				} else {
					respFlag = false
				}
			case "not contains":
				if respItem.Name == "header" && !strings.Contains(response.Header, respItem.Value) {
					respFlag = true
				} else if respItem.Name == "body" && !strings.Contains(string(response.Body), respItem.Value) {
					respFlag = true
				} else if respItem.Name == "code" && !strings.Contains(strconv.Itoa(response.Other.StatusCode), respItem.Value) {
					respFlag = true
				} else {
					respFlag = false
				}
			case "equals":
				if respItem.Name == "header" && response.Header == respItem.Value {
					respFlag = true
				} else if respItem.Name == "body" && string(response.Body) == respItem.Value {
					respFlag = true
				} else if respItem.Name == "code" && strconv.Itoa(response.Other.StatusCode) == respItem.Value {
					respFlag = true
				} else {
					respFlag = false
				}
			case "not equals":
				if respItem.Name == "header" && response.Header != respItem.Value {
					respFlag = true
				} else if respItem.Name == "body" && string(response.Body) != respItem.Value {
					respFlag = true
				} else if respItem.Name == "code" && strconv.Itoa(response.Other.StatusCode) != respItem.Value {
					respFlag = true
				} else {
					respFlag = false
				}
			case "regex":
				if resp, _ := regexp.MatchString(respItem.Value, response.Header); respItem.Name == "header" && resp {
					respFlag = true
				} else if body, _ := regexp.MatchString(respItem.Value, string(response.Body)); respItem.Name == "body" && body {
					respFlag = true
				} else if code, _ := regexp.MatchString(respItem.Value, strconv.Itoa(response.Other.StatusCode)); respItem.Name == "code" && code {
					respFlag = true
				} else {
					respFlag = false
				}
			case ">":
				if respItem.Name == "time" {
					data, err := strconv.Atoi(strings.ReplaceAll(respItem.Value, " ", ""))
					if err != nil {
						respFlag = false
					} else {
						if data > int(requestTime.Seconds()) {
							respFlag = true
						}
					}
				}
			case "<":
				if respItem.Name == "time" {
					data, err := strconv.Atoi(strings.ReplaceAll(respItem.Value, " ", ""))
					if err != nil {
						respFlag = false
					} else {
						if data < int(requestTime.Seconds()) {
							respFlag = true
						}
					}
				}
			}
			if !respFlag {
				break
			}
		}
		//----------------------------------------------
		//响应内容已检测
		//下面进入漏洞判断
		if !respFlag && vlMeta.Steps.VerifySteps.Type == "and" {
			return models.VulResult{State: false, Request: response.RequestRaw, Response: response.Header + string(response.Body)}
		}
		if respFlag && vlMeta.Steps.VerifySteps.Type == "or" {
			return models.VulResult{Request: response.RequestRaw, Response: response.Header + string(response.Body), State: true}
		}
		if respFlag && vlMeta.Steps.VerifySteps.Type == "and" {
			verifyOrFlag = true
		}
		if !respFlag && vlMeta.Steps.VerifySteps.Type == "or" {
			verifyOrFlag = false
		}
	}
	if verifyOrFlag {
		return models.VulResult{Request: response.RequestRaw, Response: response.Header + string(response.Body), State: true}
	}
	return models.VulResult{State: false, Request: response.RequestRaw, Response: response.Header + string(response.Body)}
}
